import torch
from torch.optim.lr_scheduler import ReduceLROnPlateau
from libs.loss import iou_softmax_loss, cross_entropy_fusion_loss, iou_loss
from libs.utils import save_sample_imgs, showimg, showmask, AverageMeter
from tqdm import tqdm
import importlib
from trainer.BaseTrainer import BaseTrainer
import torch.nn.functional as F
import cv2
import numpy as np

class Trainer(BaseTrainer):
    def __init__(self, config, train_data_loader, test_data_loader):
        super(Trainer, self).__init__(config, train_data_loader, test_data_loader)


    def build_model(self):
        module = importlib.import_module("nets.ESPNet")
        self.net = module.get_segmatting_model().to(self.device)


    def eval(self):
        # torch.cuda.empty_cache()
        print("star to test...")
        self.net.eval()
        with torch.no_grad():
            eval_loss_monitor = AverageMeter()
            eval_maskiouloss_monitor = AverageMeter()
            eval_alphaiouloss_monitor = AverageMeter()
            eval_loss_monitor.reset()
            eval_maskiouloss_monitor.reset()
            eval_alphaiouloss_monitor.reset()
            for samples in tqdm(self.test_data_loader):
                image = samples['image'].to(self.device)
                mask_gt = samples['mask'].to(self.device)
                alpha_gt = samples["alpha"].to(self.device)

                mask_pre, alpha_pre = self.net(image)

                evalloss, loss_alpha, loss_color, loss_ce = cross_entropy_fusion_loss(image, mask_gt, mask_pre, alpha_gt, alpha_pre, refine=self.config.refine)
                eval_loss_monitor.update(evalloss)

                mask_iou_loss = iou_softmax_loss(mask_pre, mask_gt)
                alpha_iou_loss = iou_loss(alpha_pre, alpha_gt)

                eval_maskiouloss_monitor.update(mask_iou_loss.item())
                eval_alphaiouloss_monitor.update(alpha_iou_loss.item())

            avr_loss = eval_loss_monitor.avg
            avr_maskiou = eval_maskiouloss_monitor.avg
            avr_alphaiou = eval_alphaiouloss_monitor.avg
            print("loss: %.4f mask_iou: %.4f alpha_iou: %.4f" %(avr_loss, avr_maskiou, avr_alphaiou))

        return avr_loss, avr_maskiou, avr_alphaiou



    def train(self):

        ignored_params = list(map(id, self.net.seg_extract.parameters()))
        base_params = filter(lambda p: id(p) not in ignored_params, self.net.parameters())

        if self.config.refine:
            optimizer = torch.optim.Adam([
                {'params': self.net.seg_extract.parameters(), 'lr': self.config.lr, 'weight_decay': 5e-5},
                {'params': base_params, 'lr': self.config.lr, 'weight_decay': 5e-5}
                ], lr=self.config.lr, eps=1e-7)
        else:
            optimizer = torch.optim.Adam(self.net.parameters(), lr=self.config.lr, eps=1e-7)


        lr_scheduler = ReduceLROnPlateau(optimizer, 'min', factor=0.5, patience=4, verbose=True, threshold=0.015)
        train_loss_monitor = AverageMeter()
        mask_iou_monitor = AverageMeter()
        alpha_iou_monitor = AverageMeter()

        for epoch in range(self.config.max_epoch):
            self.net.train()
            train_loss_monitor.reset()
            mask_iou_monitor.reset()
            alpha_iou_monitor.reset()

            for step, samples in enumerate(self.train_data_loader):
                image = samples['image'].to(self.device)
                mask_gt = samples['mask'].to(self.device)
                alpha_gt = samples["alpha"].to(self.device)

                mask_pre, alpha_pre = self.net(image)

                scale_mask_pre = F.softmax(mask_pre, dim=1)[0][1]

                self.net.zero_grad()

                train_loss, loss_alpha, loss_color, loss_ce = cross_entropy_fusion_loss(image, mask_gt, mask_pre, alpha_gt, alpha_pre, refine=self.config.refine)
                train_loss_monitor.update(train_loss.item())

                train_loss.backward()
                optimizer.step()

                mask_iou_loss = iou_softmax_loss(mask_pre, mask_gt)
                mask_iou_monitor.update(mask_iou_loss.item())

                alpha_iou_loss = iou_loss(alpha_pre, alpha_gt)
                alpha_iou_monitor.update(alpha_iou_loss.item())

                if step % self.config.print_step == 0:
                    print("epoch:[%d/%d] image:[%d/%d] lr:%f loss: %.4f loss_ce:%.4f mask_iou: %.4f alpha_iou: %.4f" %
                    (epoch, self.config.max_epoch, step, self.train_image_len, optimizer.param_groups[0]['lr'], train_loss.item(),
                     loss_ce.item(), mask_iou_loss.item(), alpha_iou_loss.item()))

                if 1:#step % self.config.sample_step == 0:
                    save_sample_imgs(samples, {"mask_pre": mask_pre, "alpha_pre": scale_mask_pre, "dyemask": scale_mask_pre}, self.config.sample_dir, epoch, step)


            print("average train loss: %.4f  average mask_iou: %.4f average alpha_iou: %.4f" %
                  (train_loss_monitor.avg, mask_iou_monitor.avg, alpha_iou_monitor.avg))

            #eval on Testdataset and save checkpoint
            test_loss, test_mask_iou, test_alpha_iou = self.eval()
            torch.save(self.net.state_dict(), '%s/%s_tl%0.4f_el%0.4f.pth' %
                       (self.config.checkpoint_dir, self.config.trainer, train_loss_monitor.avg, test_loss))

            lr_scheduler.step(train_loss_monitor.avg)
